import cv2
import numpy as np
import matplotlib.pyplot as plt
import math

def GenerateDCTMatrix(N):
	# http://www.lokminglui.com/dct.pdf
	# http://www.whydomath.org/node/wavlets/dct.html
	# http://www3.matapp.unimib.it/corsi-2011-2012/didattica/informatica/metodi-del-calcolo-scientifico/jpeg/papers/strang.pdf
	# http://www.mathematica-journal.com/issue/v4i1/article/81-88_Watson.mj.pdf
	result = np.zeros([N, N])
	for k in range(N):
		for x in range(N):
			if k == 0:
				constantFactor =  math.sqrt(1.0/N)
			else:
				constantFactor = math.sqrt(2.0/N)
	# optimize this to reduct compution
			result[k, x] = constantFactor*math.cos((2*x+1)*k*math.pi/(2*N))

	return result




def CompressImage(rawImg, blockSize, preserveBlockSize, showIntermediaResult = False, printMatrix = False):
	if blockSize < preserveBlockSize:
		print "The preserved block size cannot be larger than block size"
		return

	# Copy input image and reset all pixel to 0
	result = rawImg.copy()
	result.fill(0)
	# Following matrix stores intermedia result
	inteResult = np.ones(rawImg.shape, dtype=float)
	rows, cols = rawImg.shape

	# prepare DCT matrix and its transpose
	dctMat = GenerateDCTMatrix(blockSize)
	dctMatT = np.transpose(dctMat)

	# separate raw image by (blockSize*blockSize) block
	totalRowBlocks = rows/blockSize
	totalColBlocks = cols/blockSize

	#
	# Apply DCT to row image
	#
	for i in range(totalRowBlocks):
		for j in range(totalColBlocks):
			colLeftIndex = j * blockSize
			colRightIndex = (j+1) * blockSize 
			rowTopIndex = i * blockSize
			rowBottomIndex = (i+1) * blockSize

			# For Debug purpose
			# print rowTopIndex, rowBottomIndex, colLeftIndex, colRightIndex
			# print rawImg[rowTopIndex:rowBottomIndex, colLeftIndex:colRightIndex]
			# print " "
			# print np.dot(np.dot(dctMat, rawImg[rowTopIndex:rowBottomIndex, colLeftIndex:colRightIndex]), dctMatT)
			# print " "

			inteResult[rowTopIndex:rowBottomIndex, colLeftIndex:colRightIndex] = np.dot(np.dot(dctMat, rawImg[rowTopIndex:rowBottomIndex, colLeftIndex:colRightIndex]), dctMatT)

			# For Debug purpose
			# print inteResult[rowTopIndex:rowBottomIndex, colLeftIndex:colRightIndex]
			# print " "
			# print np.dot(np.dot(dctMatT, inteResult[rowTopIndex:rowBottomIndex, colLeftIndex:colRightIndex]), dctMat)
			# print " "


	if showIntermediaResult:
		cv2.namedWindow('DCT Result', cv2.CV_WINDOW_AUTOSIZE)
		cv2.imshow( "DCT Result", inteResult)

	if printMatrix:
		print inteResult


	#
	# Reduce coefficient
	#
	for i in range(totalRowBlocks):
		rowTopIndex = i * blockSize
		rowBottomIndex = (i+1) * blockSize
		inteResult[(rowTopIndex + preserveBlockSize):rowBottomIndex, :] = 0

	for j in range(totalColBlocks):
		colLeftIndex = j * blockSize
		colRightIndex = (j+1) * blockSize
		inteResult[:, (colLeftIndex + preserveBlockSize):colRightIndex] = 0

	if showIntermediaResult:
		cv2.namedWindow('Reduce coefficient Result', cv2.CV_WINDOW_AUTOSIZE)
		cv2.imshow( "Reduce coefficient Result", inteResult)


	#
	# Apply IDCT to intermedia result
	#
	for i in range(totalRowBlocks):
		for j in range(totalColBlocks):
			colLeftIndex = j * blockSize
			colRightIndex = (j+1) * blockSize
			rowTopIndex = i * blockSize
			rowBottomIndex = (i+1) * blockSize

			result[rowTopIndex:rowBottomIndex, colLeftIndex:colRightIndex] = np.dot(np.dot(dctMatT, inteResult[rowTopIndex:rowBottomIndex, colLeftIndex:colRightIndex]), dctMat)

			# For Debug purpose
			# print result[rowTopIndex:rowBottomIndex, colLeftIndex:colRightIndex]
			# print " "

	return result





if __name__ == '__main__':
	# BLOCKSIZE = 8
	# PreserverCoEfficientSize = 6
	# print "Test DCT Matrix"
	# m = GenerateDCTMatrix(BLOCKSIZE)
	# print m
	# print "=========="
	# # test data comes from http://www.whydomath.org/node/wavlets/dct.html
	# print "Test with 8*8 Matrix filled with 100"
	# testMatrix1 = np.ones([BLOCKSIZE,BLOCKSIZE])
	# testMatrix1.fill(100)

	# print testMatrix1
	# print " "

	# testMatrix1_result = np.dot(np.dot(m, testMatrix1), np.transpose(m))
	# print testMatrix1_result
	# print " "
	# print np.dot(np.dot(np.transpose(m), testMatrix1_result), m)

	# print "=========="

	# print "TEST DCT RESULT 2"
	# raw = np.ones([8, 8], dtype=float)
	# raw[0, :] = [51, 52, 51, 50, 50, 52, 50, 52]
	# raw[1, :] = [51, 52, 51, 51, 50, 52, 52, 51]
	# raw[2, :] = [50, 50, 51, 52, 52, 51, 51, 51]
	# raw[3, :] = [51, 50, 50, 50, 52, 50, 50, 51]
	# raw[4, :] = [51, 50, 50, 51, 50, 50, 51, 50]
	# raw[5, :] = [50, 51, 52, 52, 51, 50, 50, 50]
	# raw[6, :] = [51, 52, 51, 50, 52, 50, 52, 50]
	# raw[7, :] = [50, 51, 52, 52, 50, 51, 52, 51]
	# print raw

	# RESULT = np.dot(np.dot(m, raw), np.transpose(m))
	# print RESULT
	# print " "

	# # raw_Result = CompressImage(raw, 8, 6, False, True)
	# # print raw_Result
	# # print " "

	# print np.dot(np.dot(np.transpose(m), RESULT), m)

	print "========== Compress Fig3 ========== "
	rawImg = cv2.imread('./img/fig3.tif',0)
	rawImg = cv2.resize(rawImg, (rawImg.shape[0]*2, rawImg.shape[1]*2))
	cv2.namedWindow('Fig3', cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow( "Fig3", rawImg)

	compressedImg = CompressImage(rawImg, 8, 2, True)

	cv2.namedWindow('Compressed Fig3', cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow( "Compressed Fig3", compressedImg)

	cv2.namedWindow('The difference of Fig3', cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow( "The difference of Fig3", rawImg - compressedImg)

	# print "========== Compress Fig6 ========== "
	# fig6 = cv2.imread('./img/fig6.jpg',0)
	# cv2.namedWindow('Fig6', cv2.CV_WINDOW_AUTOSIZE)
	# cv2.imshow( "Fig6", fig6)

	# compressedFig6 = CompressImage(fig6, 8, 5, True)

	# cv2.namedWindow('Compressed Fig6', cv2.CV_WINDOW_AUTOSIZE)
	# cv2.imshow( "Compressed Fig6", compressedFig6)

	# cv2.namedWindow('The difference of Fig6', cv2.CV_WINDOW_AUTOSIZE)
	# cv2.imshow( "The difference of Fig6", fig6 - compressedFig6)


	print "Please check the images"


	cv2.waitKey(0)
	cv2.destroyAllWindows()